"""
MCP Client 简易实现
"""
import asyncio
import json
import os

from dotenv import load_dotenv
from fastmcp import Client
from openai import OpenAI
from openai.types.chat import ChatCompletionMessage

load_dotenv()


class UserClient:
    def __init__(self, script='server.py'):
        self.mcp_client = Client(script)  # 本地mcp服务
        self.openai_client = OpenAI(base_url=os.getenv('BASE_URL'), api_key=os.getenv('API_KEY'))  # 大模型服务
        self.model = os.getenv('MODEL')
        self.tools = []
        self.message = [
            {"role": "system", "content": "你是一个助手，请在需要的时候借助工具，使用中文回答用户问题"}
        ]

    async def prepare_tools(self):
        """
        tools = [
            {
                'type': 'function',
                'function': {
                    'name': 'multiply', 'description': 'Multiplies two numbers.',
                    'input_schema': {
                        'properties': {
                            'a': {'title': 'A', 'type': 'number'},
                            'b': {'title': 'B', 'type': 'number'}
                        },
                        'required': ['a', 'b'],
                        'type': 'object'
                    }
                }
            }
        ]
        """
        _tools = await self.mcp_client.list_tools()
        tools = [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    # "parameters": tool.parameters,
                    "input_schema": tool.inputSchema,
                },
            }
            for tool in _tools
        ]
        return tools

    def convert_args(self, func, arguments):
        """参数兼容层"""
        # arguments = json.loads(tool_call.function.arguments)
        func_conf = [i for i in self.tools if i['function']['name'] == func][0]
        converted_args = {}
        for key, value in arguments.items():
            # 可以在这里做类型转换或格式适配
            converted_args[key] = value if func_conf['function']['input_schema']['properties'][key][
                                               'type'] == 'number' else value
        print(f"Arguments: {arguments}")
        print(f"Converted arguments: {converted_args}")
        return converted_args

    async def chat(self, message: list[dict]) -> ChatCompletionMessage:
        if not self.tools:
            self.tools = await self.prepare_tools()
        response = self.openai_client.chat.completions.create(
            model=self.model,
            messages=message,
            # stream = True,
            tools=self.tools,
        )

        # 进行递归，只有当finish_reason不为tool_calls时才结束
        if response.choices[0].finish_reason != "tool_calls":
            return response.choices[0].message
        # 让client告诉大模型执行工具，多个工具需遍历执行完，并拼接所有response
        for tool_call in response.choices[0].message.tool_calls:
            func_name = tool_call.function.name
            func_args = json.loads(tool_call.function.arguments)
            # func_args = self.convert_args(func_name, func_args)
            _resp = await self.mcp_client.call_tool(func_name, func_args)
            _text = _resp.content[0].text if not _resp.is_error else f"error"
            self.message.append(
                {"role": "assistant", "content": f"{tool_call.function.name} tool call result: {_text}"}
            )
        return await self.chat(self.message)

    async def loop(self):
        async with self.mcp_client:
            await asyncio.sleep(1)
            while True:
                question = input("User: ")
                message = {"role": "user", "content": question}
                self.message.append(message)
                response_message = await self.chat(self.message)
                print("当前上下文: ", self.message)
                print(f"Bot: {response_message.content}")


async def main():
    user_client = UserClient()
    await user_client.loop()


if __name__ == '__main__':
    asyncio.run(main())
